#include <iostream>
#include "ag_robot/shnMPCfunction.h"
#include <cppad/ipopt/solve.hpp>
#include <Eigen/Core>                    // 矩阵运算
#include <Eigen/QR>

using CppAD::AD;
// FG_eval class definition implementation.
class FG_eval {
public:
    Eigen::VectorXd vec_cmd;
    Eigen::VectorXd linear_rcom;
    Eigen::VectorXd angular_rcom;
    Eigen::VectorXd robot_cmd;
    double k_h_l;
    double k_h_a;
    double _dt;
    
     double _w_linvel, _w_linvel_d, _w_angvel, _w_angvel_d;
     int _mpc_steps, _x_start, _y_start, _theta_start, _distance_start, _degree_start, _linvel_start, _angvel_start;
 
    typedef CPPAD_TESTVECTOR(AD<double>) ADvector;

    FG_eval(Eigen::VectorXd _vec_cmd , Eigen::VectorXd _linear_rcom , Eigen::VectorXd _angular_rcom,double _k_h_l,double _k_h_a)
        {
            this->vec_cmd = _vec_cmd;
            this->linear_rcom = _linear_rcom;
            this->angular_rcom = _angular_rcom;
            this->k_h_l = _k_h_l;
            this->k_h_a = _k_h_a;
            _dt = 0.25;
            _w_linvel = 5;
            _w_linvel_d = 1;
            _w_angvel = 5;
            _w_angvel_d = 1;

            _mpc_steps      = 3;
            _x_start        = 0;
            _y_start        = _x_start + _mpc_steps;
            _theta_start    = _y_start + _mpc_steps;
            _distance_start = _theta_start + _mpc_steps;
            _degree_start   = _distance_start + _mpc_steps;
            _linvel_start   = _degree_start + _mpc_steps;
            _angvel_start   = _linvel_start + _mpc_steps -1;
        } 

    FG_eval(Eigen::VectorXd _vec_cmd , Eigen::VectorXd _robot_cmd)
        {
            this->vec_cmd = _vec_cmd;
            this->robot_cmd = _robot_cmd;
            _dt = 0.25;
            _w_linvel = 5;
            _w_linvel_d = 1;
            _w_angvel = 5;
            _w_angvel_d = 1;

            _mpc_steps      = 3;
            _x_start        = 0;
            _y_start        = _x_start + _mpc_steps;
            _theta_start    = _y_start + _mpc_steps;
            _distance_start = _theta_start + _mpc_steps;
            _degree_start   = _distance_start + _mpc_steps;
            _linvel_start   = _degree_start + _mpc_steps;
            _angvel_start   = _linvel_start + _mpc_steps -1;
        }

    void LoadParams(const std::map<string, double> &params)
        {
            _dt = params.find("DT") != params.end() ? params.at("DT") : _dt;
            _mpc_steps = params.find("STEPS") != params.end() ? params.at("STEPS") : _mpc_steps;
            
            _w_linvel = params.find("W_LINVEL") != params.end() ? params.at("W_LINVEL") : _w_linvel;
            _w_linvel_d = params.find("W_DLINVEL") != params.end() ? params.at("W_DLINVEL") : _w_linvel_d;            
            _w_angvel = params.find("W_ANGVEL") != params.end() ? params.at("W_ANGVEL") : _w_angvel;
            _w_angvel_d = params.find("W_DANGVEL") != params.end() ? params.at("W_DANGVEL") : _w_angvel_d;

            //index
            _x_start        = 0;
            _y_start        = _x_start + _mpc_steps;
            _theta_start    = _y_start + _mpc_steps;
            _distance_start = _theta_start + _mpc_steps;
            _degree_start   = _distance_start + _mpc_steps;
            _linvel_start   = _degree_start + _mpc_steps;
            _angvel_start   = _linvel_start + _mpc_steps -1;
            
            //cout << "\n!! FG_eval Obj parameters updated !! " << _mpc_steps << endl; 
        } 

    void operator()(ADvector& fg, const ADvector& vars)
    {
        // //算上约束不算上下界
        // assert(fg.size() == 2);
        // assert(x.size() == 2);
        //cost function
        fg[0] = 0;
        // Minimize the use of actuators.
        for (int i = 0; i < _mpc_steps - 1 ; i++) {
            fg[0] += k_h_l * 100.0 * CppAD::pow(vars[_linvel_start + i] - vec_cmd[0], 2);
            fg[0] += k_h_a * 5.0 * CppAD::pow(vars[_angvel_start + i] - vec_cmd[1], 2);
        }
        // Minimize the use of robot.
        for (int i = 0; i < _mpc_steps - 1 ; i++) {
            fg[0] += (1-k_h_l) * 100.0 * CppAD::pow(vars[_linvel_start + i] - linear_rcom[i], 2);
            fg[0] += (1-k_h_a) * 5.0 * CppAD::pow(vars[_angvel_start + i] - angular_rcom[1], 2);
        }
        // Minimize the value gap between sequential actuations.
        for (int i = 0; i < _mpc_steps - 2; i++) {
             fg[0] += _w_linvel_d * CppAD::pow(vars[_linvel_start + i + 1] - vars[_linvel_start + i], 2);
             fg[0] += _w_angvel_d * CppAD::pow(vars[_angvel_start + i + 1] - vars[_angvel_start + i], 2);
        }

        // constraints
            fg[1 + _x_start] = vars[_x_start];
            fg[1 + _y_start] = vars[_y_start];
            fg[1 + _theta_start] = vars[_theta_start];
            fg[1 + _distance_start] = vars[_distance_start];
            fg[1 + _degree_start] = vars[_degree_start];
        //Add system dynamic model constraint
        for (int i = 0; i < _mpc_steps - 1; i++)
            {
                // The state at time t+1
                AD<double> x1 = vars[_x_start + i + 1];
                AD<double> y1 = vars[_y_start + i + 1];
                AD<double> theta1 = vars[_theta_start + i + 1];
                AD<double> distance1 = vars[_distance_start + i + 1];
                AD<double> degree1 = vars[_degree_start + i + 1];

                // The state at time t.
                AD<double> x0 = vars[_x_start + i];
                AD<double> y0 = vars[_y_start + i];
                AD<double> theta0 = vars[_theta_start + i];
                AD<double> distance0 = vars[_distance_start + i];
                AD<double> degree0 = vars[_degree_start + i];
                AD<double> degree00 = vars[_degree_start];
                AD<double> distance00 = vars[_distance_start];

                AD<double> v0 = vars[_linvel_start + i];
                AD<double> w0 = vars[_angvel_start + i];

                fg[2 + _x_start + i] = x1 - (x0 + v0 * CppAD::cos(theta0) * _dt);
                fg[2 + _y_start + i] = y1 - (y0 + v0 * CppAD::sin(theta0) * _dt);
                fg[2 + _theta_start + i] = theta1 - (theta0 +  w0 * _dt);
                // fg[2 + _distance_start + i] =distance1 - (distance00*distance00 + ((i+1)*v0*_dt)*((i+1)*v0*_dt) - 2*(distance00)*((i+1)*v0*_dt)*(CppAD::cos(degree0)));
                fg[2 + _distance_start + i] =distance1 - (distance00*distance00 + ((i+1)*0.2*_dt)*((i+1)*0.2*_dt) - 2*(distance00)*((i+1)*0.2*_dt)*(CppAD::cos(degree0)));
                fg[2 + _degree_start + i] =degree1 - (degree0 - w0 * _dt);

            }
    }

};


// ====================================
// MPC class definition implementation.
// ====================================
shnMPCFunction::shnMPCFunction() 
{
    // Set default value    
    _mpc_steps = 3;
    _max_angvel = 1.82; // Maximal angvel radian (~30 deg)
    _max_speed = 0.3; // Maximal throttle accel
    _bound_value  = 1.0e3; // Bound value for other variables

    _x_start        = 0;
    _y_start        = _x_start + _mpc_steps;
    _theta_start    = _y_start + _mpc_steps;
    _distance_start = _theta_start + _mpc_steps;
    _degree_start   = _distance_start + _mpc_steps;
    _linvel_start   = _degree_start + _mpc_steps;
    _angvel_start   = _linvel_start + _mpc_steps -1;
}

void shnMPCFunction::LoadParams(const std::map<string, double> &params)
{
    _params = params;
    //Init parameters for MPC object
    _mpc_steps = _params.find("STEPS") != _params.end() ? _params.at("STEPS") : _mpc_steps;
    _max_angvel = _params.find("MAX_ANGVEL") != _params.end() ? _params.at("MAX_ANGVEL") : _max_angvel;
    _max_speed = _params.find("MAX_SPEED") != _params.end() ? _params.at("MAX_SPEED") : _max_speed;
    _bound_value  = _params.find("BOUND") != _params.end()  ? _params.at("BOUND") : _bound_value;
    
    _x_start        = 0;
    _y_start        = _x_start + _mpc_steps;
    _theta_start    = _y_start + _mpc_steps;
    _distance_start = _theta_start + _mpc_steps;
    _degree_start   = _distance_start + _mpc_steps;
    _linvel_start   = _degree_start + _mpc_steps;
    _angvel_start   = _linvel_start + _mpc_steps -1;

    cout << "\n!! BCI MPC Obj parameters updated !! " << endl;
    // //Display the parameters
    // std::cout << "\n===== Parameters =====" << endl;
    // std::cout << "_mpc_steps: "  << _mpc_steps << endl;
    // std::cout << "_max_angvel: "  << _max_angvel << endl;
    // std::cout << "_max_speed: "  << _max_speed << endl;
    // std::cout << "_bound_value: "   << _bound_value << endl;

}

vector<double> shnMPCFunction::Solve(Eigen::VectorXd state, Eigen::VectorXd vec_cmd, Eigen::VectorXd linear_rcom,Eigen::VectorXd angular_rcom,double _speed_max, double _anguler_max,double _k_h_l,double _k_h_a) 
{   
    // std::cout << "BCI solve *******" << std::endl;
    bool ok = true;
    size_t i;
    typedef CPPAD_TESTVECTOR(double) Dvector;
    const double x = state[0];
    const double y = state[1];
    const double theta = state[2];
    const double obs_distance = state[3];
    const double obs_degree = state[4];

    // Set the number of model variables (includes both states and inputs).
    // For example: If the state is a 4 element vector, the actuators is a 2
    // element vector and there are 10 timesteps. The number of variables is:

    size_t n_vars = _mpc_steps * 5 + (_mpc_steps-1) * 2;
    
    // Set the number of constraints
    size_t n_constraints = _mpc_steps * 5;

    // Initial value of the independent variables.
    // SHOULD BE 0 besides initial state.
    Dvector vars(n_vars);
    for (int i = 0; i < n_vars; i++) 
    {
        vars[i] = 0;
    }
    
    // Set the initial variable values
    vars[_x_start] = x;
    vars[_y_start] = y;
    vars[_theta_start] = theta;
    vars[_distance_start] = obs_distance;
    vars[_degree_start] = obs_degree;

    // Set lower and upper limits for variables.
    Dvector vars_lowerbound(n_vars);
    Dvector vars_upperbound(n_vars);
    
    // Set all non-actuators upper and lowerlimits
    // to the max negative and positive values.
    for (int i = 0; i < _distance_start; i++) 
    {
        vars_lowerbound[i] = -_bound_value;
        vars_upperbound[i] = _bound_value;
    }
    // Set safty distance
    for (int i = _distance_start; i < _degree_start; i++) 
    {
        vars_lowerbound[i] = 0.5;
        vars_upperbound[i] = _bound_value;
    }

    for (int i = _degree_start; i < _linvel_start; i++) 
    {
        vars_lowerbound[i] = -_bound_value;
        vars_upperbound[i] = _bound_value;
    }

    for (int i = _linvel_start; i < _angvel_start; i++) 
    {
        vars_lowerbound[i] = 0.0;
        vars_upperbound[i] = _max_speed;
    }    

    // The upper and lower limits of angvel are set to -25 and 25
    // degrees (values in radians).
    for (int i = _angvel_start; i < n_vars; i++) 
    {
        vars_lowerbound[i] = -_max_angvel;
        vars_upperbound[i] = _max_angvel;
    }





    // Lower and upper limits for the constraints
    // Should be 0 besides initial state.
    Dvector constraints_lowerbound(n_constraints);
    Dvector constraints_upperbound(n_constraints);
    for (int i = 0; i < n_constraints; i++)
    {
        constraints_lowerbound[i] = 0;
        constraints_upperbound[i] = 0;
    }


    constraints_lowerbound[_x_start] = x;
    constraints_lowerbound[_y_start] = y;
    constraints_lowerbound[_theta_start] = theta;
    constraints_lowerbound[_distance_start] = obs_distance;
    constraints_lowerbound[_degree_start] = obs_degree;

    constraints_upperbound[_x_start] = x;
    constraints_upperbound[_y_start] = y;
    constraints_upperbound[_theta_start] = theta;
    constraints_upperbound[_distance_start] = obs_distance;
    constraints_upperbound[_degree_start] = obs_degree;

    // object that computes objective and constraints
    FG_eval fg_eval(vec_cmd,linear_rcom, angular_rcom, _k_h_l,_k_h_a);
    fg_eval.LoadParams(_params);

    // fg_eval.cost_cte;

    // options for IPOPT solver
    std::string options;
    // Uncomment this if you'd like more print information
    options += "Integer print_level  0\n";
    // NOTE: Setting sparse to true allows the solver to take advantage
    // of sparse routines, this makes the computation MUCH FASTER. If you
    // can uncomment 1 of these and see if it makes a difference or not but
    // if you uncomment both the computation time should go up in orders of
    // magnitude.
    options += "Sparse  true        forward\n";
    options += "Sparse  true        reverse\n";
    // NOTE: Currently the solver has a maximum time limit of 0.5 seconds.
    // Change this as you see fit.
    options += "Numeric max_cpu_time          0.5\n";

    // place to return solution
    CppAD::ipopt::solve_result<Dvector> solution;

    // solve the problem
    CppAD::ipopt::solve<Dvector, FG_eval>(
      options, vars, vars_lowerbound, vars_upperbound, constraints_lowerbound,
      constraints_upperbound, fg_eval, solution);

    // Check some of the solution values
    ok &= solution.status == CppAD::ipopt::solve_result<Dvector>::success;

    // Cost
    auto cost = solution.obj_value;
    // std::cout << "------------ Total Cost(solution): " << cost << "------------" << std::endl;
    // cout << "-----------------------------------------------" <<endl;
    _mpc_totalcost = cost;
    // _mpc_ctecost = Value(fg_eval.cost_cte);
    // _mpc_ethetacost = Value(fg_eval.cost_etheta);
    // _mpc_velcost = Value(fg_eval.cost_vel);

    // predicted value
    this->mpc_x = {};
    this->mpc_y = {};
    this->mpc_theta = {};
    for (int i = 0; i < _mpc_steps; i++) 
    {
        this->mpc_x.push_back(solution.x[_x_start + i]);
        this->mpc_y.push_back(solution.x[_y_start + i]);
        this->mpc_theta.push_back(solution.x[_theta_start + i]);
        //cout << "_degree_start : " << solution.x[_degree_start + i] <<endl;
        //cout << "_degree_start : " << solution.x[_distance_start + i] <<endl;
    }
    
    vector<double> result;
    result.push_back(solution.x[_linvel_start]);
    result.push_back(solution.x[_angvel_start]);
    return result;
}


vector<double> shnMPCFunction::Solve(Eigen::VectorXd state, Eigen::VectorXd vec_cmd, Eigen::VectorXd robot_cmd,double _speed_max, double _anguler_max) 
{   
    // std::cout << "BCI solve *******" << std::endl;
    bool ok = true;
    size_t i;
    typedef CPPAD_TESTVECTOR(double) Dvector;
    const double x = state[0];
    const double y = state[1];
    const double theta = state[2];
    const double obs_distance = state[3];
    const double obs_degree = state[4];

    // Set the number of model variables (includes both states and inputs).
    // For example: If the state is a 4 element vector, the actuators is a 2
    // element vector and there are 10 timesteps. The number of variables is:

    size_t n_vars = _mpc_steps * 5 + (_mpc_steps-1) * 2;
    
    // Set the number of constraints
    size_t n_constraints = _mpc_steps * 5;

    // Initial value of the independent variables.
    // SHOULD BE 0 besides initial state.
    Dvector vars(n_vars);
    for (int i = 0; i < n_vars; i++) 
    {
        vars[i] = 0;
    }
    
    // Set the initial variable values
    vars[_x_start] = x;
    vars[_y_start] = y;
    vars[_theta_start] = theta;
    vars[_distance_start] = obs_distance;
    vars[_degree_start] = obs_degree;

    // Set lower and upper limits for variables.
    Dvector vars_lowerbound(n_vars);
    Dvector vars_upperbound(n_vars);
    
    // Set all non-actuators upper and lowerlimits
    // to the max negative and positive values.
    for (int i = 0; i < _distance_start; i++) 
    {
        vars_lowerbound[i] = -_bound_value;
        vars_upperbound[i] = _bound_value;
    }
    // Set safty distance
    for (int i = _distance_start; i < _degree_start; i++) 
    {
        vars_lowerbound[i] = 0.65;
        vars_upperbound[i] = _bound_value;
    }

    for (int i = _degree_start; i < _linvel_start; i++) 
    {
        vars_lowerbound[i] = -_bound_value;
        vars_upperbound[i] = _bound_value;
    }

    for (int i = _linvel_start; i < _angvel_start; i++) 
    {
        vars_lowerbound[i] = 0.0;
        vars_upperbound[i] = _max_speed;
    }    

    // The upper and lower limits of angvel are set to -25 and 25
    // degrees (values in radians).
    for (int i = _angvel_start; i < n_vars; i++) 
    {
        vars_lowerbound[i] = -_max_angvel;
        vars_upperbound[i] = _max_angvel;
    }





    // Lower and upper limits for the constraints
    // Should be 0 besides initial state.
    Dvector constraints_lowerbound(n_constraints);
    Dvector constraints_upperbound(n_constraints);
    for (int i = 0; i < n_constraints; i++)
    {
        constraints_lowerbound[i] = 0;
        constraints_upperbound[i] = 0;
    }


    constraints_lowerbound[_x_start] = x;
    constraints_lowerbound[_y_start] = y;
    constraints_lowerbound[_theta_start] = theta;
    constraints_lowerbound[_distance_start] = obs_distance;
    constraints_lowerbound[_degree_start] = obs_degree;

    constraints_upperbound[_x_start] = x;
    constraints_upperbound[_y_start] = y;
    constraints_upperbound[_theta_start] = theta;
    constraints_upperbound[_distance_start] = obs_distance;
    constraints_upperbound[_degree_start] = obs_degree;

    // object that computes objective and constraints
    FG_eval fg_eval(vec_cmd,robot_cmd);
    fg_eval.LoadParams(_params);

    // fg_eval.cost_cte;

    // options for IPOPT solver
    std::string options;
    // Uncomment this if you'd like more print information
    options += "Integer print_level  0\n";
    // NOTE: Setting sparse to true allows the solver to take advantage
    // of sparse routines, this makes the computation MUCH FASTER. If you
    // can uncomment 1 of these and see if it makes a difference or not but
    // if you uncomment both the computation time should go up in orders of
    // magnitude.
    options += "Sparse  true        forward\n";
    options += "Sparse  true        reverse\n";
    // NOTE: Currently the solver has a maximum time limit of 0.5 seconds.
    // Change this as you see fit.
    options += "Numeric max_cpu_time          0.5\n";

    // place to return solution
    CppAD::ipopt::solve_result<Dvector> solution;

    // solve the problem
    CppAD::ipopt::solve<Dvector, FG_eval>(
      options, vars, vars_lowerbound, vars_upperbound, constraints_lowerbound,
      constraints_upperbound, fg_eval, solution);

    // Check some of the solution values
    ok &= solution.status == CppAD::ipopt::solve_result<Dvector>::success;

    // Cost
    auto cost = solution.obj_value;
    // std::cout << "------------ Total Cost(solution): " << cost << "------------" << std::endl;
    // cout << "-----------------------------------------------" <<endl;
    _mpc_totalcost = cost;
    // _mpc_ctecost = Value(fg_eval.cost_cte);
    // _mpc_ethetacost = Value(fg_eval.cost_etheta);
    // _mpc_velcost = Value(fg_eval.cost_vel);

    // predicted value
    this->mpc_x = {};
    this->mpc_y = {};
    this->mpc_theta = {};
    for (int i = 0; i < _mpc_steps; i++) 
    {
        this->mpc_x.push_back(solution.x[_x_start + i]);
        this->mpc_y.push_back(solution.x[_y_start + i]);
        this->mpc_theta.push_back(solution.x[_theta_start + i]);
        //cout << "_degree_start : " << solution.x[_degree_start + i] <<endl;
        //cout << "_degree_start : " << solution.x[_distance_start + i] <<endl;
    }
    
    vector<double> result;
    result.push_back(solution.x[_linvel_start]);
    result.push_back(solution.x[_angvel_start]);
    return result;
}
